import os
import re
import graphviz

# 解析SQL文件（有中文注释）
def parse_sql_file(file_path):
    with open(file_path, 'r', encoding='utf-8') as file:
        content = file.read()

    # 匹配表的注释和字段的注释
    table_pattern = re.compile(r"CREATE TABLE `([^`]+)`[\s\S]+?COMMENT = '(.*?)' ROW_FORMAT = D", re.S)
    column_pattern = re.compile(r"`([^`]+)`\s+([^\s(]+(?:\([^)]+\))?)\s+.*?COMMENT '(.*?)'", re.S)
    primary_key_pattern = re.compile(r"PRIMARY KEY \(([^)]+)\)", re.S)  # 用于提取主键字段列表


    tables = {}

    # 查找所有表的信息
    for table_match in table_pattern.finditer(content):
        table_name = table_match.group(1)
        table_comment = table_match.group(2)

        # print(f"正在解析表：{table_name} - {table_comment}")  # 调试信息，查看表是否正确匹配

        # 查找当前表的字段信息
        columns = []
        table_content = table_match.group(0)

        # 查找主键字段
        primary_key_columns = []
        primary_key_match = primary_key_pattern.search(table_content)
        if primary_key_match:
            # 提取多个主键字段并清理空格和反引号
            primary_key_columns = [col.strip(' `') for col in primary_key_match.group(1).split(',')]
            # print(f"  主键字段：{primary_key_columns}")  # 调试信息，查看主键字段是否正确匹配

        for column_match in column_pattern.finditer(table_content):
            column_name = column_match.group(1)
            column_type = column_match.group(2)
            column_comment = column_match.group(3)

            # print(f"    字段名：{column_name}, 字段类型：{column_type}, 字段注释：{column_comment}")  # 添加调试信息

            # 清理字段名（去掉反引号）
            column_name = column_name.strip(' `')

            # 检查字段是否是主键
            is_primary_key = column_name in primary_key_columns

            columns.append({
                "name": column_name,
                "type": column_type,
                "comment": column_comment,
                "is_primary_key": is_primary_key
            })

        # 添加表的信息到字典
        tables[table_name] = {
            "comment": table_comment,
            "columns": columns
        }

    return tables

# 解析SQL文件（无注释或部分注释）
def parse_sql_file_without_comments(file_path):
        with open(file_path, 'r', encoding='utf-8') as file:
            content = file.read()

        # 匹配表的注释和字段的注释
        # table_pattern = re.compile(r"CREATE TABLE `([^`]+)`[\s\S]+?AUTO_INCREMENT = (.*?) +ROW_FORMAT = D", re.S)
        table_pattern = re.compile(r"CREATE TABLE `([^`]+)`[\s\S]+?ENGINE = (.*?) +ROW_FORMAT = D", re.S)

        column_pattern = re.compile(r"`([^`]+)`\s+([^\s(]+(?:\([^)]+\))?)\s+.*?COMMENT '(.*?)'", re.S)
        primary_key_pattern = re.compile(r"PRIMARY KEY \(([^)]+)\)", re.S)  # 用于提取主键字段列表

        table_comment_pattern = re.compile(r"COMMENT = '(.*?)'", re.S)
        column_comment_pattern = re.compile(r"COMMENT '(.*?)'", re.S)

        tables = {}

        # 查找所有表的信息
        for table_match in table_pattern.finditer(content):
            table_name = table_match.group(1)
            have_table_comment = table_comment_pattern.search(table_match.group(0))
            table_comment = have_table_comment.group(1) if have_table_comment else table_name
            # table_comment =

            # print(f"正在解析表：{table_name} - {table_comment}")  # 调试信息，查看表是否正确匹配

            # 查找当前表的字段信息
            columns = []
            table_content = table_match.group(0)

            # 查找主键字段
            primary_key_columns = []
            primary_key_match = primary_key_pattern.search(table_content)
            if primary_key_match:
                # 提取多个主键字段并清理空格和反引号
                primary_key_columns = [col.strip(' `') for col in primary_key_match.group(1).split(',')]
                # print(f"  主键字段：{primary_key_columns}")  # 调试信息，查看主键字段是否正确匹配

            for column_match in column_pattern.finditer(table_content):
                column_name = column_match.group(1)
                column_type = column_match.group(2)
                have_column_comment = column_comment_pattern.search(column_match.group(0))
                column_comment = have_column_comment.group(1) if have_column_comment else column_name

                # print(f"    字段名：{column_name}, 字段类型：{column_type}, 字段注释：{column_comment}")  # 添加调试信息

                # 清理字段名（去掉反引号）
                column_name = column_name.strip(' `')

                # 检查字段是否是主键
                is_primary_key = column_name in primary_key_columns

                columns.append({
                    "name": column_name,
                    "type": column_type,
                    "comment": column_comment,
                    "is_primary_key": is_primary_key
                })

            # 添加表的信息到字典
            tables[table_name] = {
                "comment": table_comment,
                "columns": columns
            }

        return tables

# 画ER图（单实体）
def generate_er_diagram(tables, output_dir, er_chinese):
    """
    根据解析到的表结构信息生成ER图。
    """
    for table_name, table_info in tables.items():

        table_comment = table_info['comment'] if er_chinese else table_name
        columns = table_info['columns']

        # 创建Graphviz图表
        dot = graphviz.Digraph(
            name=table_name,
            format='png',
            engine='circo',
            graph_attr={'labeljust': 'c', 'margin': '0.5'},
            node_attr={
                'style': 'filled',
                'fillcolor': 'white',
                'fontname': 'SimHei',
                'fontcolor': 'black'
            }
        )

        # 添加实体节点（矩形），使用中文注释作为标签
        dot.node('ENTITY', label=table_comment or table_name, shape='rectangle', width='1.5', height='1')

        # 添加属性节点（椭圆），并对主键添加下划线标注
        for column in columns:
            label = f"{column['comment']}" if er_chinese else column['name']
            if column['is_primary_key']:
                label = f"_{label}_"  # 主键加下划线标注
            dot.node(column['name'], label=label, shape='ellipse')

        # 添加不可见边控制布局，并设置箭头为无
        for column in columns:
            dot.edge('ENTITY', column['name'], arrowhead='none', color='black')

        # 生成文件
        output_path = os.path.join(output_dir, f'{table_name}')
        dot.render(filename=output_path, cleanup=True)
        print(f"表 {table_name} 的单实体 ER 图已成功渲染为 {table_name}.png")

    print("文件已输出到'./export/'目录")
    print("完成！")

